import argparse
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import os


def handleData(dataPath):
    with open(dataPath ,'r') as f:
        rawData = f.readlines()
    edges = np.zeros((eval(rawData[1].split()[-2]),eval(rawData[1].split()[-2])) ,np.int32 )
    numberOfEdges = eval(rawData[1].split()[1])
    for edgeNO in range(numberOfEdges):
        node1, node2 = eval(rawData[edgeNO + 2].split()[0]) - 1, eval(rawData[edgeNO + 2].split()[1]) - 1
        edges[node1, node2] = 1
        edges[node2, node1] = 1
    return edges

def drawGraph( edges, nodePos = {1:(9,-1)} ): 
    G = nx.Graph()
    G.add_nodes_from([i + 1 for i in range(len(edges))  ])
    for i in range(len(edges)):
        for j in range(len(edges)):
            if edges[i,j] == 1:
                G.add_edge(i+1,1+j)
    if nodePos[1][1] == -1:
        nx.draw_networkx(G)
    else:
        nx.draw_networkx(G,nodePos)
    plt.axis('off')
    os.makedirs('./pic/',exist_ok = True)
    plt.savefig('./pic/origin.png',dpi = 400)

def drawResult( communities, edges): 
    G = nx.Graph()
    G.add_nodes_from([i + 1 for i in range(len(edges))  ])
    nodeColors = []
    nodePos = dict()
    for i in range(len(edges)):
        if (i+1) in communities[0]:
            nodeColors.append('blue')
            nodePos[i+1] = (np.random.rand()*8 - 2 ,np.random.rand()*15) 
        else:
            nodeColors.append('green')
            nodePos[i+1] = (np.random.rand()*8 + 5.5 ,np.random.rand()*15) 

        for j in range(len(edges)):
            if edges[i,j] == 1:
                G.add_edge(i+1,1+j)            
    nx.draw_networkx(G,pos=nodePos, node_color = nodeColors)
    plt.axis('off')
    os.makedirs('./pic/',exist_ok = True)
    plt.savefig('./pic/community.png',dpi = 400)
    return nodePos


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_path',default='./ucidata-zachary/out.ucidata-zachary',type=str,required=False,help='数据集路径')

    args = parser.parse_args()
    dataPath = args.data_path
    edges = handleData(dataPath)
    drawGraph(edges)
    print(edges)

if __name__  == '__main__':
    main()
